import torch
import torch.nn as nn
import torchvision.transforms as transforms
import PIL
import pandas as pd
from PIL import Image
import numpy as np
import os
from torch.utils.data import DataLoader

class DRB(nn.Module):
    def __init__(self,in_channels,out_channels,h,w):
        super(DRB,self).__init__()
        self.in_channels=in_channels
        self.out_channels=out_channels

        self.conv1=nn.Conv2d(self.in_channels,64,3,1,1)

        self.conv2=nn.Conv2d(64,128,1,1,0)
        self.relu1=nn.ReLU(inplace=True)
        self.conv3=nn.Conv2d(128,256,1,1,0)
        self.conv3_1 = nn.Conv2d(256, 64, 1, 1, 0)
        self.global_pool1=nn.AdaptiveAvgPool2d((h,w))
        self.sigmoid1=nn.Sigmoid()

        self.relu2 = nn.ReLU(inplace=True)
        self.conv4=nn.Conv2d(64,64,3,1,1)
        self.global_pool2=nn.AdaptiveAvgPool2d((h,w))
        self.sigmoid2=nn.Sigmoid()


    def forward(self,x):

        residual=x

        out1=self.conv1(x)

        out2=self.conv2(out1)
        out2=self.relu1(out2)
        out2=self.conv3_1(self.conv3(out2))

        out2=out2+0.2*self.sigmoid1(self.global_pool1(out2))

        out3=out1+out2

        out4=self.relu2(out3)
        out4=self.conv4(out4)
        out4=out4+self.sigmoid2(self.global_pool2(out4))

        out=out4+residual
        return out

class DFMU(nn.Module):
    def __init__(self,h,w):
        super(DFMU,self).__init__()
        self.RG1 = self._make_layers(DRB,8,h,w)
        self.RG2 = self._make_layers(DRB,8,h,w)
        self.RG3 = self._make_layers(DRB,8,h,w)
    def _make_layers(self,block,block_nums,h,w):
        layers=[]
        for _ in range(0,block_nums):
            layers.append(block(64,64,h,w))
        return nn.Sequential(*layers)


    def forward(self,x):

        out1= self.RG1(x)
        out2= self.RG2(out1)
        out3= self.RG3(out2)

        return 0.4*out1,0.3*out2,0.2*out3


class SFMU(nn.Module):
    def __init__(self):
        super(SFMU,self).__init__()
        self.conv=nn.Conv2d(64,64,1,1,0)

        self.conv1=nn.Conv2d(64,64,1,1,0)
        self.conv3=nn.Conv2d(64,64,3,1,1)
        self.conv5=nn.Conv2d(64,64,5,1,2)

    def forward(self,x):
        x=self.conv(x)

        out1=self.conv1(x)

        out2=0.2*out1+x
        out2=self.conv3(out2)

        out3=0.2*out2+x
        out3=self.conv3(out3)

        return 0.8*out1,0.7*out2,0.8*out3


class MFFN(nn.Module):
    def __init__(self,in_channels,out_channels,h,w):
        super(MFFN,self).__init__()
        self.in_channels=in_channels
        self.out_channels=out_channels
        self.conv=nn.Conv2d(self.in_channels,64,3,1,1)
        self.dfmu=DFMU(h,w)
        self.sfmu=SFMU()

        self.conv1=nn.Conv2d(128,128,1,1,0)
        self.pixelShuffler1=nn.PixelShuffle(4)

        self.conv2 = nn.Conv2d(128, 128, 1, 1, 0)
        self.pixelShuffler2 = nn.PixelShuffle(4)

        self.conv3 = nn.Conv2d(128, 128, 1, 1, 0)
        self.pixelShuffler3 = nn.PixelShuffle(4)

        self.last_conv=nn.Conv2d(24,self.out_channels,3,1,1)



    def forward(self,x):
        x=self.conv(x)

        p1,p2,p3=self.dfmu(x)
        s1,s2,s3=self.sfmu(x)

        I1 = self.pixelShuffler1(self.conv1(torch.cat((p1,s1),dim=1)))
        I2 = self.pixelShuffler2(self.conv2(torch.cat((p2, s2), dim=1)))
        I3 = self.pixelShuffler3(self.conv3(torch.cat((p3, s3), dim=1)))

        out=self.last_conv(torch.cat((I1,I2,I3),dim=1))
        return out










if __name__=="__main__":
    input=torch.randn((1,3,64,64))
    mmfn=MFFN(3,3,64,64)
    f3=mmfn(input)


    print(f3.size())




